import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import numpy as np

from plotly.offline import init_notebook_mode
init_notebook_mode()

# Load and preprocess data
df = pd.read_csv('../data/capture_fisheries.csv', skiprows=4)
df = df.dropna(how='all')

# Clean data
df['Country Name'] = df['Country Name'].str.replace('"', '')
years = [str(year) for year in range(1960, 2025)]
id_vars = ["Country Name", "Country Code", "Indicator Name", "Indicator Code"]

# Melt data
df_melted = df.melt(id_vars=id_vars, value_vars=years, 
                   var_name="Year", value_name="Production")
df_melted['Year'] = pd.to_numeric(df_melted['Year'])
df_melted['Production'] = pd.to_numeric(df_melted['Production'], errors='coerce')
df_melted = df_melted.dropna(subset=['Production'])

# Calculate focused range (exclude extreme outliers)
q95 = df_melted['Production'].quantile(0.95)
focused_data = df_melted[df_melted['Production'] <= q95]
zmax = focused_data['Production'].max()
zmin = focused_data['Production'].min()

# Enhanced color scale with more distinction
enhanced_scale = [
    [0.0, '#E1F5FE'],  # Darkest purple
    [0.1, '#B3E5FC'],  # Deep purple
    [0.2, '#81D4FA'],  # Purple
    [0.3, '#4FC3F7'],  # Magenta
    [0.4, '#29B6F6'],  # Pink
    [0.5, '#039BE5'],  # Red
    [0.6, '#0288D1'],  # Orange-red
    [0.7, '#0277BD'],  # Orange
    [0.8, '#01579B'],  # Light orange
    [0.9, '#0D47A1'],  # Pale yellow
    [1.0, '#1d1c47']   # White
]

# 1. Static map showing total production (with focused scale)
total_production = df_melted.groupby(['Country Name', 'Country Code'])['Production'].sum().reset_index()

fig_total = px.choropleth(total_production,
                        locations="Country Code",
                        color="Production",
                        hover_name="Country Name",
                        color_continuous_scale=enhanced_scale,
                        range_color=[zmin, zmax],
                        title="<b>Total Fisheries Production per country</b>",
                        labels={'Production': 'Production (metric tons)'},
                        projection="natural earth")

fig_total.update_layout(
    geo=dict(
        landcolor='#f0f0f0',
        lakecolor='#d0e0f0',
        bgcolor='white',
        showframe=True,
        framecolor='black'
    ),
    coloraxis_colorbar=dict(
        title="Production",
        thickness=20,
        len=0.75,
        tickvals=np.linspace(zmin, zmax, 6),
        ticktext=[f"{int(x):,}" for x in np.linspace(zmin, zmax, 6)]
    ),
    margin={"r":0,"t":60,"l":0,"b":0}
)

# 2. Interactive map with year slider - with better color distribution
frames = []
for year in sorted(df_melted['Year'].unique()):
    yearly_data = df_melted[df_melted['Year'] == year]
    yearly_focused = yearly_data[yearly_data['Production'] <= yearly_data['Production'].quantile(0.95)]
    year_zmax = yearly_focused['Production'].max()
    year_zmin = yearly_focused['Production'].min()
    
    frames.append(go.Frame(
        data=[go.Choropleth(
            locations=yearly_data['Country Code'],
            z=yearly_data['Production'],
            zmin=year_zmin,
            zmax=year_zmax,
            colorscale=enhanced_scale,
            customdata=yearly_data[['Country Name']],
            hovertemplate="<b>%{customdata[0]}</b><br>" +
                         "Year: %{frame.name}<br>" +
                         "Production: %{z:,} metric tons<extra></extra>",
            marker_line_color='darkgray',
            marker_line_width=0.5
        )],
        name=str(year)
    ))

# Create base figure
fig_slider = go.Figure(
    data=[go.Choropleth(
        locations=df_melted['Country Code'],
        z=df_melted['Production'],
        colorscale=enhanced_scale,
        customdata=df_melted[['Country Name']],
        hovertemplate="<b>%{customdata[0]}</b><br>" +
                     "Year: %{frame.name}<br>" +
                     "Production: %{z:,} metric tons<extra></extra>",
        marker_line_color='darkgray',
        marker_line_width=0.5
    )],
    layout=go.Layout(
        title="<b>Annual Fisheries Production per country</b>",
        geo=dict(
            showframe=True,
            showcoastlines=True,
            landcolor='#f0f0f0',
            lakecolor='#d0e0f0',
            bgcolor='white',
            framecolor='black',
            framewidth=1
        ),
        sliders=[{
            "active": 0,
            "steps": [{
                "args": [[f.name], {"frame": {"duration": 300, "redraw": True},
                                  "mode": "immediate"}],
                "label": f.name,
                "method": "animate"
            } for f in frames],
            "x": 0.1,
            "len": 0.9,
            "currentvalue": {
                "prefix": "<b>Year: </b>",
                "font": {"size": 14},
                "xanchor": "center"
            }
        }],
        updatemenus=[{
            "buttons": [
                {
                    "args": [None, {"frame": {"duration": 500, "redraw": True},
                                  "fromcurrent": True}],
                    "label": "▶️ Play",
                    "method": "animate"
                },
                {
                    "args": [[None], {"frame": {"duration": 0, "redraw": True},
                                    "mode": "immediate"}],
                    "label": "⏸️ Pause",
                    "method": "animate"
                }
            ],
            "type": "buttons",
            "x": 0.1,
            "y": 0,
            "bgcolor": "rgba(255,255,255,0.8)",
            "bordercolor": "#444"
        }]
    ),
    frames=frames
)

# Add explanatory annotation
fig_slider.add_annotation(
    x=0.5,
    y=-0.1,
    xref="paper",
    yref="paper",
    text="Color scale automatically focuses on 95% of data range for each year",
    showarrow=False,
    font=dict(size=12, color="#444")
)

fig_total.show()
fig_slider.show()